#!/usr/bin/python
# -*-coding:utf-8-*-
'''
拼接数据的API接口，用于拼接因子数据和标签数据
'''
import pandas as pd
import os

# try:
#     try:
#         from support.factors_library_base import DataReader
#     except:
#         from .support.factors_library_base import DataReader
# except:
#     from zg_data_process.support.factors_library_base import DataReader

# from zg_data_process.support.dir_info import buffer_data_dir
from zbc_factor_lib.dir_info import cache_data_dir as buffer_data_dir

# from zg02_factor_lib.base.factors_library_base import NewFactorLib as DataReader
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

'''
'stock_code',
'date',
'board',
'down_limit',
'down_one_line',
'is_csi300',
'is_csi500',
'is_sh50',
'paused',
'sw1_code',
'sw1_name',
'up_limit',
'up_one_line',
'is_st',
'stock_name',
'list_date',
'delist_date',
'list_day',
'is_new',
'is_sub_new'
'''

db = 'validation'

# data_api = DataReader(db=db)
# data_api.show_factor_library_db()

class DataConcat(DataReader):
    def __init__(self,
                 label_data_filename='processed_new_stock_label_ci1_data',
                 size_db=db):
        super(DataConcat, self).__init__()
        # self.factor_name_list = ['scale_total_market_size']
        self.factor_name_list = ['scale_circulate_market_size']
        self.label_data_filename = label_data_filename
        self.db = size_db
        self.buffer_data_dir = buffer_data_dir

    def cache_data(self, df, filename, format='table', main_columns=None):
        '''
        缓存数据

        :param df: pandas Dataframe
        :param filename:
        :param main_columns:
        :return:
        '''
        df.to_hdf(os.path.join(self.buffer_data_dir, '%s.h5' % (filename)),
                  key=filename,
                  format=format,
                  data_columns=main_columns,
                  mode='w')

    def read_cache_data(self, filename, filter_list=None, columns=[], start_row=None, end_row=None):
        '''
        获取缓存数据

        :param filename:
        :param filter_list:
        :param columns:
        :param start_row:
        :param end_row:
        :return:
        '''
        if len(columns) == 0:
            df = pd.read_hdf(os.path.join(self.buffer_data_dir, '%s.h5' % (filename)),
                             key=filename,
                             mode='r',
                             start=start_row,
                             end=end_row,
                             where=filter_list)
        else:
            df = pd.read_hdf(os.path.join(self.buffer_data_dir, '%s.h5' % (filename)),
                             key=filename,
                             mode='r',
                             start=start_row,
                             end=end_row,
                             columns=columns,
                             where=filter_list)
        return df

    def input(self, factor_name_list):
        self.factor_name_list.extend(factor_name_list)
        self.factor_name_list = list(set(self.factor_name_list))

    def label_factor_data(self,
                          df,
                          start_date=None,
                          end_date=None,
                          cache=False,
                          cache_filename=None,
                          add_size_factor=True,
                          refresh=True,
                          verbose=True,
                          copy=True):
        '''
        用于给输入的因子数据打标签

        输入df - pandas dataframe 不带index 需要stock_code和date两个索引指标在columns里面
        '''
        if verbose: print('label given factor data starts...')

        if cache_filename is not None:
            # 查看是否有cache数据
            if os.path.exists(os.path.join(self.buffer_data_dir, cache_filename+'.h5')) and not refresh:
                concat_data = self.read_cache_data(filename=cache_filename)
                if verbose: print('cache data loaded!\n')
                return concat_data

        if start_date is None and end_date is not None:
            filter_list = ["date <= '%s'"  % (end_date)]
        elif start_date is not None and end_date is None:
            filter_list = ["date >= '%s'"  % (start_date)]
        elif start_date is not None and end_date is not None:
            filter_list = ["date >= '%s' & date <= '%s'"  % (start_date, end_date)]
        else:
            filter_list = None

        # # TODO - 读取标签数据
        # stock_label_data = self.read_basic_data_table(filename=self.label_data_filename,
        #                                         filter_list=None)
        #                                         # filter_list=filter_list)
        #
        # stock_label_data = stock_label_data[(stock_label_data['date'] >= start_date) &
        #                                     (stock_label_data['date'] <= end_date)]
        #
        # stock_label_data = stock_label_data.set_index(['stock_code', 'date'])
        # if verbose: print('stock label data loaded!')

        # TODO - 读取标签数据
        stock_label_data = self.read_basic_data_table(filename=self.label_data_filename,
                                                      filter_list=None)
                                                # filter_list=filter_list)

        if start_date is None and end_date is not None:
            stock_label_data = stock_label_data[stock_label_data['date'] <= end_date]
        elif start_date is not None and end_date is None:
            stock_label_data = stock_label_data[stock_label_data['date'] >= start_date]
        elif start_date is not None and end_date is not None:
            stock_label_data = stock_label_data[(stock_label_data['date'] >= start_date) &
                                                (stock_label_data['date'] <= end_date)]
        else:
            pass

        stock_label_data = stock_label_data.set_index(['stock_code', 'date'])
        if verbose: print('stock label data loaded!')

        # TODO - 因子数据并入
        if copy:
            df = df.copy()

        df = df.set_index(['stock_code', 'date'])

        # TODO - 合并数据
        stock_label_data = stock_label_data.reindex(df.index)

        stock_label_data = stock_label_data.loc[df.index, :]

        if add_size_factor:
            # TODO - 市值数据并入
            # size_factor_data = self.read_factor_table(filename='scale_total_market_size',
            # size_factor_data = self.read_factor_table(filename='scale_circulate_market_size',
            size_factor_data = self.read_factor_table(filename='circulate_cap',
                                                      filter_list=filter_list)

            size_factor_data = size_factor_data.set_index(['stock_code', 'date'])

            size_factor_data = size_factor_data.rename(
                # columns={'scale_circulate_market_size': 'scale_total_market_size'}
                columns={'circulate_cap': 'scale_total_market_size'}
            )

            if verbose: print('total size data loaded!')

            size_factor_data = size_factor_data.reindex(stock_label_data.index)

            concat_data = pd.concat([stock_label_data, size_factor_data, df], axis=1)
                                    # join_axes=[stock_label_data.index])
        else:
            concat_data = pd.concat([stock_label_data, df], axis=1)
                                    # join_axes=[stock_label_data.index])

        concat_data.index.names = ['stock_code', 'date']
        concat_data = concat_data.reset_index()

        # TODO - cache 数据
        if cache:
            assert cache_filename is not None

            self.cache_data(df=concat_data,
                            filename=cache_filename,
                            main_columns=['stock_code', 'date'])

            if verbose: print('concat data cached!')

        if verbose: print('data concated!\n')

        return concat_data

    def label_selected_factor_data(self,
                                   df,
                                   start_date=None,
                                   end_date=None,
                                   cache=False,
                                   selected_factor=None,
                                   cache_filename=None,
                                   add_size_factor=True,
                                   refresh=True,
                                   verbose=True,
                                   copy=True):
        '''
        用于给输入的因子数据打标签，只标注行业和市值信息

        输入df - pandas dataframe 不带index 需要stock_code和date两个索引指标在columns里面
        '''
        if verbose: print('label given factor data starts...')

        if cache_filename is not None:
            # 查看是否有cache数据
            if os.path.exists(os.path.join(self.buffer_data_dir, cache_filename+'.h5')) and not refresh:
                concat_data = self.read_cache_data(filename=cache_filename)

                if verbose: print('cache data loaded!\n')

                return concat_data

        if start_date is None and end_date is not None:
            filter_list = ["date <= '%s'"  % (end_date)]
        elif start_date is not None and end_date is None:
            filter_list = ["date >= '%s'"  % (start_date)]
        elif start_date is not None and end_date is not None:
            filter_list = ["date >= '%s' & date <= '%s'"  % (start_date, end_date)]
        else:
            filter_list = None

        # TODO - 读取标签数据
        if selected_factor is not None:
            stock_label_data = self.read_basic_data_table(filename=self.label_data_filename,
                                                        filter_list=filter_list,
                                                        columns=['stock_code', 'date'] + selected_factor)
        else:
            stock_label_data = self.read_basic_data_table(filename=self.label_data_filename,
                                                        filter_list=filter_list,
                                                        columns=['stock_code', 'date', 'sw1_code'])

        stock_label_data = stock_label_data.set_index(['stock_code', 'date'])
        if verbose: print('stock label data loaded!')

        # TODO - 因子数据并入
        if copy:
            df = df.copy()
        df = df.set_index(['stock_code', 'date'])

        # TODO - 合并数据
        stock_label_data = stock_label_data.loc[df.index, :]

        if add_size_factor:
            # TODO - 市值数据并入
            # size_factor_data = self.read_factor_table(filename='scale_total_market_size',
            size_factor_data = self.read_factor_table(filename='circulate_cap',
                                                      filter_list=filter_list)

            size_factor_data = size_factor_data.set_index(['stock_code', 'date'])

            size_factor_data = size_factor_data.rename(
                # columns={'scale_circulate_market_size': 'scale_total_market_size'}
                columns={'circulate_cap': 'scale_total_market_size'}
            )

            if verbose: print('total size data loaded!')

            size_factor_data = size_factor_data.reindex(stock_label_data.index)

            concat_data = pd.concat([stock_label_data, size_factor_data, df], axis=1)
                                    # join_axes=[stock_label_data.index])
        else:
            concat_data = pd.concat([stock_label_data, df], axis=1)
                                    # join_axes=[stock_label_data.index])

        concat_data.index.names = ['stock_code', 'date']
        concat_data = concat_data.reset_index()

        # TODO - cache 数据
        if cache:
            assert cache_filename is not None
            self.cache_data(df=concat_data,
                            filename=cache_filename,
                            main_columns=['stock_code', 'date'])
            if verbose: print('concat data cached!')
        if verbose: print('data concated!\n')

        return concat_data

    def label_cap_factor_data(self,
                               df,
                               start_date=None,
                               end_date=None,
                               cache=False,
                               cache_filename=None,
                               refresh=True,
                               verbose=True,
                               copy=True):
        '''
        用于给输入的因子数据打标签，只标注市值信息

        输入df - pandas dataframe 不带index 需要stock_code和date两个索引指标在columns里面
        '''
        if copy:
            df = df.copy()

        if verbose:
            print('label given factor data starts...')

        if cache_filename is not None:
            # 查看是否有cache数据
            if os.path.exists(os.path.join(self.buffer_data_dir, cache_filename+'.h5')) and not refresh:
                concat_data = self.read_cache_data(filename=cache_filename)
                if verbose: print('cache data loaded!\n')
                return concat_data

        if start_date is None and end_date is not None:
            filter_list = ["date <= '%s'"  % (end_date)]
        elif start_date is not None and end_date is None:
            filter_list = ["date >= '%s'"  % (start_date)]
        elif start_date is not None and end_date is not None:
            filter_list = ["date >= '%s' & date <= '%s'"  % (start_date, end_date)]
        else:
            filter_list = None

        # TODO - 市值数据并入
        # size_factor_data = self.read_factor_table(filename='scale_total_market_size',
        # size_factor_data = self.read_factor_table(filename='scale_circulate_market_size',
        size_factor_data = self.read_factor_table(filename='circulate_cap',
                                                  filter_list=filter_list)

        size_factor_data = size_factor_data.set_index(['stock_code', 'date'])

        size_factor_data = size_factor_data.rename(
            # columns={'scale_circulate_market_size': 'scale_total_market_size'}
            columns={'circulate_cap': 'scale_total_market_size'}
        )

        # TODO - 合并数据
        code_date_multi_index = df[['stock_code', 'date']].set_index(['stock_code', 'date']).index

        concat_data = df.copy()
        concat_data['scale_total_market_size'] = size_factor_data.loc[code_date_multi_index, :].values
        del code_date_multi_index

        # TODO - cache 数据
        if cache:
            assert cache_filename is not None
            self.cache_data(df=concat_data,
                            filename=cache_filename,
                            main_columns=['stock_code', 'date'])

            if verbose:
                print('concat data cached!')

        if verbose:
            print('cap data concated!\n')

        return concat_data

    def label_factor_data_with_given_label_data(self,
                                                 df,
                                                 stock_label_data=None,
                                                 size_factor_data=None,
                                                 start_date=None,
                                                 end_date=None,
                                                 cache=False,
                                                 cache_filename=None,
                                                 add_size_factor=True,
                                                 refresh=True,
                                                 verbose=True,
                                                 copy=True):
        '''
        用于给输入的因子数据打标签

        输入df - pandas dataframe 不带index 需要stock_code和date两个索引指标在columns里面
        '''
        if verbose: print('label given factor data starts...')

        if cache_filename is not None:
            # 查看是否有cache数据
            if os.path.exists(os.path.join(self.buffer_data_dir, cache_filename+'.h5')) and not refresh:
                concat_data = self.read_cache_data(filename=cache_filename)

                if verbose: print('cache data loaded!\n')

                return concat_data

        # if start_date is None and end_date is not None:
        #     filter_list = ["date <= '%s'"  % (end_date)]
        # elif start_date is not None and end_date is None:
        #     filter_list = ["date >= '%s'"  % (start_date)]
        # elif start_date is not None and end_date is not None:
        #     filter_list = ["date >= '%s' & date <= '%s'"  % (start_date, end_date)]
        # else:
        #     filter_list = None

        # # TODO - 读取标签数据
        # stock_label_data = self.read_basic_data_table(filename=self.label_data_filename,
        #                                         filter_list=None)
        #                                         # filter_list=filter_list)
        #
        # stock_label_data = stock_label_data[(stock_label_data['date'] >= start_date) &
        #                                     (stock_label_data['date'] <= end_date)]
        #
        # stock_label_data = stock_label_data.set_index(['stock_code', 'date'])
        # if verbose: print('stock label data loaded!')

        if stock_label_data is None:
            # TODO - 读取标签数据
            stock_label_data = self.read_basic_data_table(filename=self.label_data_filename,
                                                          filter_list=None)
        else:
            stock_label_data = stock_label_data.copy()

        if start_date is None and end_date is not None:
            stock_label_data = stock_label_data[stock_label_data['date'] <= end_date]
        elif start_date is not None and end_date is None:
            stock_label_data = stock_label_data[stock_label_data['date'] >= start_date]
        elif start_date is not None and end_date is not None:
            stock_label_data = stock_label_data[(stock_label_data['date'] >= start_date) &
                                                (stock_label_data['date'] <= end_date)]
        else:
            pass

        stock_label_data = stock_label_data.set_index(['stock_code', 'date'])
        if verbose: print('stock label data loaded!')

        # TODO - 因子数据并入
        if copy:
            df = df.copy()

        df = df.set_index(['stock_code', 'date'])

        # TODO - 合并数据
        stock_label_data = stock_label_data.loc[df.index, :]

        if add_size_factor:
            # TODO - 市值数据并入
            if size_factor_data is None:
                # size_factor_data = self.read_factor_table(filename='scale_total_market_size',
                # size_factor_data = self.read_factor_table(filename='scale_circulate_market_size',
                size_factor_data = self.read_factor_table(filename='circulate_cap',
                                                          filter_list=None)
            else:
                size_factor_data = size_factor_data.copy()

            if start_date is None and end_date is not None:
                size_factor_data = size_factor_data[size_factor_data['date'] <= end_date]
            elif start_date is not None and end_date is None:
                size_factor_data = size_factor_data[size_factor_data['date'] >= start_date]
            elif start_date is not None and end_date is not None:
                size_factor_data = size_factor_data[(size_factor_data['date'] >= start_date) &
                                                    (size_factor_data['date'] <= end_date)]
            else:
                pass


            size_factor_data = size_factor_data.set_index(['stock_code', 'date'])

            size_factor_data = size_factor_data.rename(
                # columns={'scale_circulate_market_size': 'scale_total_market_size'}
                columns={'circulate_cap': 'scale_total_market_size'}
            )

            if verbose: print('total size data loaded!')

            size_factor_data = size_factor_data.reindex(stock_label_data.index)

            concat_data = pd.concat([stock_label_data, size_factor_data, df], axis=1)
                                    # join_axes=[stock_label_data.index])
        else:
            concat_data = pd.concat([stock_label_data, df], axis=1)
                                    # join_axes=[stock_label_data.index])

        concat_data.index.names = ['stock_code', 'date']
        concat_data = concat_data.reset_index()

        # TODO - cache 数据
        if cache:
            assert cache_filename is not None
            self.cache_data(df=concat_data,
                            filename=cache_filename,
                            main_columns=['stock_code', 'date'])
            if verbose: print('concat data cached!')
        if verbose: print('data concated!\n')

        return concat_data


if __name__ == '__main__':
    data_concat_api = DataConcat(label_data_filename='process_stock_label_sw1_data')

    data_concat_api.show_factor_library_db()

    # # parameters
    # factor_name_list = ['tech_return_1_predictability_12', 'profitability_roa_ttm']
    # start_date = '2010-01-01'
    # end_date = '2018-07-31'
    #
    # # concat data
    # data_concat_api.input(factor_name_list)
    # concat_data = data_concat_api.concatenate(start_date=start_date,
    #                                           end_date=end_date,
    #                                           cache=True,
    #                                           cache_filename='zbc_concat_cache_data',
    #                                           refresh=True,
    #                                           verbose=True)

    buffer_dir = '/db/zg_data/zbc/buffer'
    test_data = pd.read_hdf(os.path.join(buffer_dir, 'zbc_monthly_factor_data.h5'))
    start = test_data['date'].min().strftime('%Y-%m-%d')
    end = test_data['date'].max().strftime('%Y-%m-%d')

    from time import time


    st = time()
    labeled_factor_data = data_concat_api.label_selected_factor_data(df=test_data,
                                                                     start_date=start,
                                                                     end_date=end,
                                                                     cache=True,
                                                                     add_size_factor=True,
                                                                     cache_filename='zbc_test_labeled_cache_data',
                                                                     refresh=True,
                                                                     verbose=True)
    et = time()
    print('add one label with cap time is %.5f sec.' % (et-st))

    st = time()
    labeled_factor_data = data_concat_api.label_selected_factor_data(df=test_data,
                                                                     start_date=start,
                                                                     end_date=end,
                                                                     cache=True,
                                                                     add_size_factor=False,
                                                                     cache_filename='zbc_test_labeled_cache_data',
                                                                     refresh=True,
                                                                     verbose=True)
    et = time()
    print('add one label without cap time is %.5f sec.' % (et-st))

    st = time()
    labeled_factor_data = data_concat_api.label_factor_data(df=test_data,
                                                            start_date=start,
                                                            end_date=end,
                                                            cache=True,
                                                            add_size_factor=True,
                                                            cache_filename='zbc_test_labeled_cache_data',
                                                            refresh=True,
                                                            verbose=True)
    et = time()
    print('add all labels time is %.5f sec.' % (et-st))

    # # check
    # labeled_factor_data.head()
    # multi_index = concat_data.index
    # unique_multi_index = multi_index.unique()
    # print(multi_index.shape)
    # print(unique_multi_index.shape)
